2923. Дерево

Задано подвешенное дерево, состоящее из n вершин. Каждая вершина окрашена в один из n цветов. Для каждой вершины v требуется определить количество различных цветов, встречающихся в поддереве с корнем в v.

Вход. В первой строке задано число n (1 ≤ n ≤ 10⁶). Далее следуют n строк, каждая из которых описывает одну вершину. В i-й строке указаны два числа p_i c_i, где p_i – номер родителя вершины i, а c_i – цвет вершины i (1 ≤ c_i ≤ n). Для корня дерева p_i = 0.

Выход. Выведите n чисел – по одному для каждой вершины от 1 до n. Для каждой вершины укажите количество различных цветов в поддереве с корнем в ней.

Пример входа	Пример выхода
5 2 1 3 2 0 3 3 3 2 1	1 2 3 1 1

РЕШЕНИЕ

поиск в глубину

Анализ алгоритма

Выполним обход дерева в глубину, начиная с корня. Для каждой вершины i будем хранить множество s_i, в котором накапливаются цвета всех вершин, входящих в поддерево с корнем в i.

Если j – сын вершины i в процессе обхода, то множество s_j должно быть включено в s_i.

Количество различных цветов в поддереве с корнем в вершине i равно мощности (размеру) множества s_i.

Пример

Слева для каждой вершины указан её цвет, справа – множество цветов, содержащихся в поддереве с корнем в этой вершине.

Реализация алгоритма

В массиве color[i] хранится цвет вершины i. Во множестве s[i] будем накапливать цвета, встречающиеся в поддереве с корнем в вершине i.

Ориентированное дерево представлено в виде списка смежности графа g. В массиве res[i] сохраняется количество различных цветов в поддеревье с корнем в вершине i.

#define MAX 1000010

int color[MAX], res[MAX];

set<int> s[MAX];

vector<vector<int> > g;

Функция dfs выполняет обход дерева в глубину, начиная с вершины v.

void dfs(int v)

{

Сначала в множество s[v] добавляем цвет самой вершины v.

s[v].insert(color[v]);

Перебираем ребра дерева (v, to).

for (int to : g[v])

{

dfs(to);

Для каждого ребра (v, to) добавляем множество s[to] в s[v]. Если размер множества s[v] меньше размера множества s[to], то меняем их местами. Далее содержимое меньшего множества s[to] добавляем во множество s[v].

if (s[v].size() < s[to].size()) s[v].swap(s[to]);

s[v].insert(s[to].begin(), s[to].end());

Очищаем множество s[to] – оно нам больше не пригодится.

s[to].clear();

}

После этого в res[v] записываем количество различных цветов в поддереве, то есть размер множества s[v].

res[v] = s[v].size();

}

Основная часть программы. Читаем входные данные.

scanf("%d",&n);

g.resize(n+1);

for(i = 1; i <= n; i++)

{

scanf("%d %d",&p,&c);

g[p].push_back(i);

color[i] = c;

}

Запускаем поиск в глубину из корня дерева – вершины с номером 0.

dfs(0);

Выводим ответ.

for(i = 1; i <= n; i++)

printf("%d ",res[i]);

printf("\n");